/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package mysqlctl

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"reflect"
	"sort"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/mysql/fakesqldb"
	"vitess.io/vitess/go/mysql/replication"
	"vitess.io/vitess/go/protoutil"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/test/utils"
	"vitess.io/vitess/go/vt/logutil"
	"vitess.io/vitess/go/vt/mysqlctl/backupstats"
	"vitess.io/vitess/go/vt/mysqlctl/backupstorage"

	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

const mysqlShutdownTimeout = 1 * time.Minute

// TestBackupExecutesBackupWithScopedParams tests that Backup passes
// a Scope()-ed stats to backupengine ExecuteBackup.
func TestBackupExecutesBackupWithScopedParams(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)

	require.Nil(t, Backup(env.ctx, env.backupParams), env.logger.Events)

	require.Equal(t, 1, len(env.backupEngine.ExecuteBackupCalls))
	executeBackupParams := env.backupEngine.ExecuteBackupCalls[0].BackupParams
	var executeBackupStats *backupstats.FakeStats
	for _, sr := range env.stats.ScopeReturns {
		if sr == executeBackupParams.Stats {
			executeBackupStats = sr
		}
	}
	require.Contains(t, executeBackupStats.ScopeV, backupstats.ScopeComponent)
	require.Equal(t, backupstats.BackupEngine.String(), executeBackupStats.ScopeV[backupstats.ScopeComponent])
	require.Contains(t, executeBackupStats.ScopeV, backupstats.ScopeImplementation)
	require.Equal(t, "Fake", executeBackupStats.ScopeV[backupstats.ScopeImplementation])
}

// TestBackupNoStats tests that if BackupParams.Stats is nil, then Backup will
// pass non-nil Stats to sub-components.
func TestBackupNoStats(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)
	env.setStats(nil)

	require.Nil(t, Backup(env.ctx, env.backupParams), env.logger.Events)

	// It parameterizes the backup storage with nop stats.
	require.Equal(t, 1, len(env.backupStorage.WithParamsCalls))
	require.Equal(t, backupstats.NoStats(), env.backupStorage.WithParamsCalls[0].Stats)
}

// TestBackupParameterizesBackupStorageWithScopedStats tests that Backup passes
// a Scope()-ed stats to BackupStorage.WithParams.
func TestBackupParameterizesBackupStorageWithScopedStats(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)

	require.Nil(t, Backup(env.ctx, env.backupParams), env.logger.Events)

	require.Equal(t, 1, len(env.backupStorage.WithParamsCalls))
	var storageStats *backupstats.FakeStats
	for _, sr := range env.stats.ScopeReturns {
		if sr == env.backupStorage.WithParamsCalls[0].Stats {
			storageStats = sr
		}
	}
	require.Contains(t, storageStats.ScopeV, backupstats.ScopeComponent)
	require.Equal(t, backupstats.BackupStorage.String(), storageStats.ScopeV[backupstats.ScopeComponent])
	require.Contains(t, storageStats.ScopeV, backupstats.ScopeImplementation)
	require.Equal(t, "Fake", storageStats.ScopeV[backupstats.ScopeImplementation])
}

// TestBackupEmitsStats tests that Backup emits stats.
func TestBackupEmitsStats(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)

	// Force ExecuteBackup to take time so we can test stats emission.
	env.backupEngine.ExecuteBackupDuration = 1001 * time.Millisecond

	require.Nil(t, Backup(env.ctx, env.backupParams), env.logger.Events)

	require.NotZero(t, backupstats.DeprecatedBackupDurationS.Get())
	require.Equal(t, 0, len(env.stats.TimedIncrementCalls))
	require.Equal(t, 0, len(env.stats.ScopeV))
}

// TestBackupTriesToParameterizeBackupStorage tests that Backup tries to pass
// backupstorage.Params to backupstorage, but only if it responds to
// backupstorage.WithParams.
func TestBackupTriesToParameterizeBackupStorage(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)

	require.Nil(t, Backup(env.ctx, env.backupParams), env.logger.Events)

	require.Equal(t, 1, len(env.backupStorage.WithParamsCalls))
	require.Equal(t, env.logger, env.backupStorage.WithParamsCalls[0].Logger)
	var scopedStats backupstats.Stats
	for _, sr := range env.stats.ScopeReturns {
		if sr != env.backupStorage.WithParamsCalls[0].Stats {
			continue
		}
		if scopedStats != nil {
			require.Fail(t, "backupstorage stats matches multiple scoped stats produced by parent stats")
		}
		scopedStats = sr
	}
	require.NotNil(t, scopedStats)
}

func TestFindFilesToBackupWithoutRedoLog(t *testing.T) {
	root := t.TempDir()

	// Initialize the fake mysql root directories
	innodbDataDir := path.Join(root, "innodb_data")
	innodbLogDir := path.Join(root, "innodb_log")
	dataDir := path.Join(root, "data")
	dataDbDir := path.Join(dataDir, "vt_db")
	extraDir := path.Join(dataDir, "extra_dir")
	outsideDbDir := path.Join(root, "outside_db")
	rocksdbDir := path.Join(dataDir, ".rocksdb")
	sdiOnlyDir := path.Join(dataDir, "sdi_dir")
	for _, s := range []string{innodbDataDir, innodbLogDir, dataDbDir, extraDir, outsideDbDir, rocksdbDir, sdiOnlyDir} {
		err := os.MkdirAll(s, os.ModePerm)
		require.NoErrorf(t, err, "failed to create directory %v: %v", s, err)
	}

	innodbLogFile := "innodb_log_1"

	if err := os.WriteFile(path.Join(innodbDataDir, "innodb_data_1"), []byte("innodb data 1 contents"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file innodb_data_1: %v", err)
	}
	if err := os.WriteFile(path.Join(innodbLogDir, innodbLogFile), []byte("innodb log 1 contents"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file %s: %v", innodbLogFile, err)
	}
	if err := os.WriteFile(path.Join(dataDbDir, "db.opt"), []byte("db opt file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file db.opt: %v", err)
	}
	if err := os.WriteFile(path.Join(extraDir, "extra.stuff"), []byte("extra file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file extra.stuff: %v", err)
	}
	if err := os.WriteFile(path.Join(outsideDbDir, "table1.frm"), []byte("frm file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file table1.opt: %v", err)
	}
	if err := os.Symlink(outsideDbDir, path.Join(dataDir, "vt_symlink")); err != nil {
		t.Fatalf("failed to symlink vt_symlink: %v", err)
	}
	if err := os.WriteFile(path.Join(rocksdbDir, "000011.sst"), []byte("rocksdb file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file 000011.sst: %v", err)
	}
	if err := os.WriteFile(path.Join(sdiOnlyDir, "table1.sdi"), []byte("sdi file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file table1.sdi: %v", err)
	}

	cnf := &Mycnf{
		InnodbDataHomeDir:     innodbDataDir,
		InnodbLogGroupHomeDir: innodbLogDir,
		DataDir:               dataDir,
	}

	result, totalSize, err := findFilesToBackup(cnf)
	if err != nil {
		t.Fatalf("findFilesToBackup failed: %v", err)
	}
	sort.Sort(forTest(result))
	t.Logf("findFilesToBackup returned: %v", result)
	expected := []FileEntry{
		{
			Base: "Data",
			Name: ".rocksdb/000011.sst",
		},
		{
			Base: "Data",
			Name: "sdi_dir/table1.sdi",
		},
		{
			Base: "Data",
			Name: "vt_db/db.opt",
		},
		{
			Base: "Data",
			Name: "vt_symlink/table1.frm",
		},
		{
			Base: "InnoDBData",
			Name: "innodb_data_1",
		},
		{
			Base: "InnoDBLog",
			Name: innodbLogFile,
		},
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("got wrong list of FileEntry %v, expected %v", result, expected)
	}
	if totalSize <= 0 {
		t.Fatalf("backup size should be > 0, got %v", totalSize)
	}
}

func TestFindFilesToBackupWithRedoLog(t *testing.T) {
	root := t.TempDir()

	// Initialize the fake mysql root directories
	innodbDataDir := path.Join(root, "innodb_data")
	innodbLogDir := path.Join(root, "innodb_log")
	dataDir := path.Join(root, "data")
	dataDbDir := path.Join(dataDir, "vt_db")
	extraDir := path.Join(dataDir, "extra_dir")
	outsideDbDir := path.Join(root, "outside_db")
	rocksdbDir := path.Join(dataDir, ".rocksdb")
	sdiOnlyDir := path.Join(dataDir, "sdi_dir")
	for _, s := range []string{innodbDataDir, innodbLogDir, dataDbDir, extraDir, outsideDbDir, rocksdbDir, sdiOnlyDir} {
		if err := os.MkdirAll(s, os.ModePerm); err != nil {
			t.Fatalf("failed to create directory %v: %v", s, err)
		}
	}

	cnf := &Mycnf{
		InnodbDataHomeDir:     innodbDataDir,
		InnodbLogGroupHomeDir: innodbLogDir,
		DataDir:               dataDir,
	}

	os.Mkdir(path.Join(innodbLogDir, mysql.DynamicRedoLogSubdir), os.ModePerm)
	innodbLogFile := path.Join(mysql.DynamicRedoLogSubdir, "#ib_redo1")

	if err := os.WriteFile(path.Join(innodbDataDir, "innodb_data_1"), []byte("innodb data 1 contents"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file innodb_data_1: %v", err)
	}
	if err := os.WriteFile(path.Join(innodbLogDir, innodbLogFile), []byte("innodb log 1 contents"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file %s: %v", innodbLogFile, err)
	}
	if err := os.WriteFile(path.Join(dataDbDir, "db.opt"), []byte("db opt file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file db.opt: %v", err)
	}
	if err := os.WriteFile(path.Join(extraDir, "extra.stuff"), []byte("extra file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file extra.stuff: %v", err)
	}
	if err := os.WriteFile(path.Join(outsideDbDir, "table1.frm"), []byte("frm file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file table1.opt: %v", err)
	}
	if err := os.Symlink(outsideDbDir, path.Join(dataDir, "vt_symlink")); err != nil {
		t.Fatalf("failed to symlink vt_symlink: %v", err)
	}
	if err := os.WriteFile(path.Join(rocksdbDir, "000011.sst"), []byte("rocksdb file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file 000011.sst: %v", err)
	}
	if err := os.WriteFile(path.Join(sdiOnlyDir, "table1.sdi"), []byte("sdi file"), os.ModePerm); err != nil {
		t.Fatalf("failed to write file table1.sdi: %v", err)
	}

	result, totalSize, err := findFilesToBackup(cnf)
	if err != nil {
		t.Fatalf("findFilesToBackup failed: %v", err)
	}
	sort.Sort(forTest(result))
	t.Logf("findFilesToBackup returned: %v", result)
	expected := []FileEntry{
		{
			Base: "Data",
			Name: ".rocksdb/000011.sst",
		},
		{
			Base: "Data",
			Name: "sdi_dir/table1.sdi",
		},
		{
			Base: "Data",
			Name: "vt_db/db.opt",
		},
		{
			Base: "Data",
			Name: "vt_symlink/table1.frm",
		},
		{
			Base: "InnoDBData",
			Name: "innodb_data_1",
		},
		{
			Base: "InnoDBLog",
			Name: innodbLogFile,
		},
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("got wrong list of FileEntry %v, expected %v", result, expected)
	}
	if totalSize <= 0 {
		t.Fatalf("backup size should be > 0, got %v", totalSize)
	}
}

// TestRestoreEmitsStats tests that Restore emits stats.
func TestRestoreEmitsStats(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)

	// Force ExecuteRestore to take time so we can test stats emission.
	env.backupEngine.ExecuteRestoreDuration = 1001 * time.Millisecond

	_, err := Restore(env.ctx, env.restoreParams)
	require.Nil(t, err, env.logger.Events)

	require.NotZero(t, backupstats.DeprecatedRestoreDurationS.Get())
	require.Equal(t, 0, len(env.stats.TimedIncrementCalls))
	require.Equal(t, 0, len(env.stats.ScopeV))
}

// TestRestoreExecutesRestoreWithScopedParams tests that Restore passes
// a Scope()-ed stats to backupengine ExecuteRestore.
func TestRestoreExecutesRestoreWithScopedParams(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)

	_, err := Restore(env.ctx, env.restoreParams)
	require.Nil(t, err, env.logger.Events)

	require.Equal(t, 1, len(env.backupEngine.ExecuteRestoreCalls))
	executeRestoreParams := env.backupEngine.ExecuteRestoreCalls[0].RestoreParams
	var executeRestoreStats *backupstats.FakeStats
	for _, sr := range env.stats.ScopeReturns {
		if sr == executeRestoreParams.Stats {
			executeRestoreStats = sr
		}
	}
	require.Contains(t, executeRestoreStats.ScopeV, backupstats.ScopeComponent)
	require.Equal(t, backupstats.BackupEngine.String(), executeRestoreStats.ScopeV[backupstats.ScopeComponent])
	require.Contains(t, executeRestoreStats.ScopeV, backupstats.ScopeImplementation)
	require.Equal(t, "Fake", executeRestoreStats.ScopeV[backupstats.ScopeImplementation])
}

// TestRestoreNoStats tests that if RestoreParams.Stats is nil, then Restore will
// pass non-nil Stats to sub-components.
func TestRestoreNoStats(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)
	env.setStats(nil)

	_, err := Restore(env.ctx, env.restoreParams)
	require.Nil(t, err, env.logger.Events)

	// It parameterizes the backup storage with nop stats.
	require.Equal(t, 1, len(env.backupStorage.WithParamsCalls))
	require.Equal(t, backupstats.NoStats(), env.backupStorage.WithParamsCalls[0].Stats)
}

// TestRestoreParameterizesBackupStorageWithScopedStats tests that Restore passes
// a Scope()-ed stats to BackupStorage.WithParams.
func TestRestoreParameterizesBackupStorageWithScopedStats(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)

	_, err := Restore(env.ctx, env.restoreParams)
	require.Nil(t, err, env.logger.Events)

	require.Equal(t, 1, len(env.backupStorage.WithParamsCalls))
	var storageStats *backupstats.FakeStats
	for _, sr := range env.stats.ScopeReturns {
		if sr == env.backupStorage.WithParamsCalls[0].Stats {
			storageStats = sr
		}
	}
	require.Contains(t, storageStats.ScopeV, backupstats.ScopeComponent)
	require.Equal(t, backupstats.BackupStorage.String(), storageStats.ScopeV[backupstats.ScopeComponent])
	require.Contains(t, storageStats.ScopeV, backupstats.ScopeImplementation)
	require.Equal(t, "Fake", storageStats.ScopeV[backupstats.ScopeImplementation])
}

// TestRestoreTriesToParameterizeBackupStorage tests that Restore tries to pass
// backupstorage.Params to backupstorage, but only if it responds to
// backupstorage.WithParams.
func TestRestoreTriesToParameterizeBackupStorage(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)

	_, err := Restore(env.ctx, env.restoreParams)
	require.Nil(t, err, env.logger.Events)

	require.Equal(t, 1, len(env.backupStorage.WithParamsCalls))
	require.Equal(t, env.logger, env.backupStorage.WithParamsCalls[0].Logger)
	var scopedStats backupstats.Stats
	for _, sr := range env.stats.ScopeReturns {
		if sr != env.backupStorage.WithParamsCalls[0].Stats {
			continue
		}
		if scopedStats != nil {
			require.Fail(t, "backupstorage stats matches multiple scoped stats produced by parent stats")
		}
		scopedStats = sr
	}
	require.NotNil(t, scopedStats)
}

// TestRestoreManifestMySQLVersionValidation tests that Restore tries to validate
// the MySQL version and safe upgrade attribute.
func TestRestoreManifestMySQLVersionValidation(t *testing.T) {
	testCases := []struct {
		fromVersion, toVersion string
		upgradeSafe            bool
		wantErr                bool
	}{
		{
			fromVersion: "mysqld  Ver 5.7.42",
			toVersion:   "mysqld  Ver 8.0.32",
			upgradeSafe: true,
			wantErr:     false,
		},
		{
			fromVersion: "mysqld  Ver 5.7.42",
			toVersion:   "mysqld  Ver 8.0.32",
			upgradeSafe: false,
			wantErr:     true,
		},
		{
			fromVersion: "mysqld  Ver 5.7.42",
			toVersion:   "mysqld  Ver 8.0.32",
			upgradeSafe: true,
			wantErr:     false,
		},
		{
			fromVersion: "mysqld  Ver 8.0.32",
			toVersion:   "mysqld  Ver 8.0.32",
			upgradeSafe: false,
			wantErr:     false,
		},
		{
			fromVersion: "mysqld  Ver 8.0.32",
			toVersion:   "mysqld  Ver 8.0.32",
			upgradeSafe: true,
			wantErr:     false,
		},
		{
			fromVersion: "mysqld  Ver 8.0.32",
			toVersion:   "mysqld  Ver 8.0.31",
			upgradeSafe: false,
			wantErr:     true,
		},
		{
			fromVersion: "mysqld  Ver 8.0.32",
			toVersion:   "mysqld  Ver 8.0.31",
			upgradeSafe: true,
			wantErr:     true,
		},
		{
			fromVersion: "mysqld  Ver 8.0.32",
			toVersion:   "mysqld  Ver 8.0.33",
			upgradeSafe: false,
			wantErr:     true,
		},
		{
			fromVersion: "mysqld  Ver 8.0.32",
			toVersion:   "mysqld  Ver 8.0.33",
			upgradeSafe: true,
			wantErr:     false,
		},
		{
			fromVersion: "",
			toVersion:   "mysqld  Ver 8.0.33",
			upgradeSafe: false,
			wantErr:     false,
		},
		{
			fromVersion: "",
			toVersion:   "mysqld  Ver 8.0.33",
			upgradeSafe: true,
			wantErr:     false,
		},
	}

	for _, tc := range testCases {
		t.Run(fmt.Sprintf("%s->%s upgradeSafe=%t", tc.fromVersion, tc.toVersion, tc.upgradeSafe), func(t *testing.T) {
			env := createFakeBackupRestoreEnv(t)
			env.mysqld.Version = tc.toVersion

			manifest := BackupManifest{
				BackupTime:   time.Now().Add(-1 * time.Hour).Format(time.RFC3339),
				BackupMethod: fakeBackupEngineName,
				Keyspace:     "test",
				Shard:        "-",
				MySQLVersion: tc.fromVersion,
				UpgradeSafe:  tc.upgradeSafe,
			}

			manifestBytes, err := json.Marshal(manifest)
			require.Nil(t, err)

			env.backupEngine.ExecuteRestoreReturn = FakeBackupEngineExecuteRestoreReturn{&manifest, nil}
			env.backupStorage.ListBackupsReturn = FakeBackupStorageListBackupsReturn{
				BackupHandles: []backupstorage.BackupHandle{
					&FakeBackupHandle{
						ReadFileReturnF: func(context.Context, string) (io.ReadCloser, error) {
							return io.NopCloser(bytes.NewBuffer(manifestBytes)), nil
						},
					},
				},
			}

			_, err = Restore(env.ctx, env.restoreParams)
			if tc.wantErr {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
			}
		})
	}
}

type forTest []FileEntry

func (f forTest) Len() int           { return len(f) }
func (f forTest) Swap(i, j int)      { f[i], f[j] = f[j], f[i] }
func (f forTest) Less(i, j int) bool { return f[i].Base+f[i].Name < f[j].Base+f[j].Name }

type fakeBackupRestoreEnv struct {
	backupEngine  *FakeBackupEngine
	backupParams  BackupParams
	backupStorage *FakeBackupStorage
	ctx           context.Context
	logger        *logutil.MemoryLogger
	restoreParams RestoreParams
	mysqld        *FakeMysqlDaemon
	stats         *backupstats.FakeStats
}

func createFakeBackupRestoreEnv(t *testing.T) *fakeBackupRestoreEnv {
	ctx := context.Background()
	logger := logutil.NewMemoryLogger()

	sqldb := fakesqldb.New(t)
	sqldb.SetNeverFail(true)
	mysqld := NewFakeMysqlDaemon(sqldb)
	require.Nil(t, mysqld.Shutdown(ctx, nil, false, mysqlShutdownTimeout))

	dirName, err := os.MkdirTemp("", "vt_backup_test")
	require.Nil(t, err)

	cnf := &Mycnf{
		DataDir: dirName,
	}

	stats := backupstats.NewFakeStats()

	backupParams := BackupParams{
		Cnf:                  cnf,
		Logger:               logger,
		Mysqld:               mysqld,
		Concurrency:          1,
		HookExtraEnv:         map[string]string{},
		TopoServer:           nil,
		Keyspace:             "test",
		Shard:                "-",
		BackupTime:           time.Now(),
		IncrementalFromPos:   "",
		Stats:                stats,
		MysqlShutdownTimeout: mysqlShutdownTimeout,
	}

	restoreParams := RestoreParams{
		Cnf:                  cnf,
		Logger:               logger,
		Mysqld:               mysqld,
		Concurrency:          1,
		HookExtraEnv:         map[string]string{},
		DeleteBeforeRestore:  false,
		DbName:               "test",
		Keyspace:             "test",
		Shard:                "-",
		StartTime:            time.Now(),
		RestoreToPos:         replication.Position{},
		DryRun:               false,
		Stats:                stats,
		MysqlShutdownTimeout: mysqlShutdownTimeout,
	}

	manifest := BackupManifest{
		BackupTime:   FormatRFC3339(time.Now().Add(-1 * time.Hour)),
		BackupMethod: fakeBackupEngineName,
		Keyspace:     "test",
		Shard:        "-",
		MySQLVersion: "8.0.32",
	}

	manifestBytes, err := json.Marshal(manifest)
	require.Nil(t, err)

	testBackupEngine := FakeBackupEngine{}
	testBackupEngine.ExecuteRestoreReturn = FakeBackupEngineExecuteRestoreReturn{&manifest, nil}

	previousBackupEngineImplementation := backupEngineImplementation
	BackupRestoreEngineMap[fakeBackupEngineName] = &testBackupEngine
	backupEngineImplementation = fakeBackupEngineName

	testBackupStorage := FakeBackupStorage{}
	testBackupStorage.ListBackupsReturn = FakeBackupStorageListBackupsReturn{
		BackupHandles: []backupstorage.BackupHandle{
			&FakeBackupHandle{
				ReadFileReturnF: func(context.Context, string) (io.ReadCloser, error) {
					return io.NopCloser(bytes.NewBuffer(manifestBytes)), nil
				},
			},
		},
	}
	testBackupStorage.StartBackupReturn = FakeBackupStorageStartBackupReturn{&FakeBackupHandle{}, nil}
	testBackupStorage.WithParamsReturn = &testBackupStorage

	backupstorage.BackupStorageMap[fakeBackupEngineName] = &testBackupStorage
	previousBackupStorageImplementation := backupstorage.BackupStorageImplementation
	backupstorage.BackupStorageImplementation = fakeBackupEngineName

	// all restore integration tests must be leak checked
	t.Cleanup(func() {
		utils.EnsureNoLeaks(t)
	})

	t.Cleanup(func() {
		backupstats.DeprecatedBackupDurationS.Reset()
		backupstats.DeprecatedRestoreDurationS.Reset()

		delete(BackupRestoreEngineMap, fakeBackupEngineName)
		backupEngineImplementation = previousBackupEngineImplementation

		delete(backupstorage.BackupStorageMap, fakeBackupEngineName)
		backupstorage.BackupStorageImplementation = previousBackupStorageImplementation
		mysqld.Close()
		sqldb.Close()
	})

	return &fakeBackupRestoreEnv{
		backupEngine:  &testBackupEngine,
		backupParams:  backupParams,
		backupStorage: &testBackupStorage,
		ctx:           ctx,
		logger:        logger,
		mysqld:        mysqld,
		restoreParams: restoreParams,
		stats:         stats,
	}
}

func (fbe *fakeBackupRestoreEnv) setStats(stats *backupstats.FakeStats) {
	fbe.backupParams.Stats = nil
	fbe.restoreParams.Stats = nil
	fbe.stats = nil
}

func TestParseBackupName(t *testing.T) {
	// backup name doesn't contain 3 parts
	_, _, err := ParseBackupName("dir", "asd.saddsa")
	assert.ErrorContains(t, err, "cannot backup name")

	// Invalid time
	bt, al, err := ParseBackupName("dir", "2024-03-18.123.tablet_id")
	assert.Nil(t, bt)
	assert.Nil(t, al)
	assert.NoError(t, err)

	// Valid case
	bt, al, err = ParseBackupName("dir", "2024-03-18.180911.cell1-42")
	assert.NotNil(t, *bt, time.Date(2024, 03, 18, 18, 9, 11, 0, time.UTC))
	assert.Equal(t, "cell1", al.Cell)
	assert.Equal(t, uint32(42), al.Uid)
	assert.NoError(t, err)
}

func TestShouldRestore(t *testing.T) {
	env := createFakeBackupRestoreEnv(t)

	b, err := ShouldRestore(env.ctx, env.restoreParams)
	assert.False(t, b)
	assert.Error(t, err)

	env.restoreParams.DeleteBeforeRestore = true
	b, err = ShouldRestore(env.ctx, env.restoreParams)
	assert.True(t, b)
	assert.NoError(t, err)
	env.restoreParams.DeleteBeforeRestore = false

	env.mysqld.FetchSuperQueryMap = map[string]*sqltypes.Result{
		"SHOW DATABASES": {Rows: [][]sqltypes.Value{{sqltypes.NewVarBinary("any_db")}}},
	}
	b, err = ShouldRestore(env.ctx, env.restoreParams)
	assert.NoError(t, err)
	assert.True(t, b)

	env.mysqld.FetchSuperQueryMap = map[string]*sqltypes.Result{
		"SHOW DATABASES": {Rows: [][]sqltypes.Value{{sqltypes.NewVarBinary("test")}}},
	}
	b, err = ShouldRestore(env.ctx, env.restoreParams)
	assert.False(t, b)
	assert.NoError(t, err)
}

func TestScanLinesToLogger(t *testing.T) {
	reader, writer := io.Pipe()
	logger := logutil.NewMemoryLogger()
	var wg sync.WaitGroup

	wg.Add(1)
	go scanLinesToLogger("test", reader, logger, wg.Done)

	for i := range 100 {
		_, err := fmt.Fprintf(writer, "foobar %d\n", i)
		require.NoError(t, err)
	}

	writer.Close()
	wg.Wait()

	require.Equal(t, 100, len(logger.Events))

	for i, event := range logger.Events {
		require.Equal(t, fmt.Sprintf("test: foobar %d", i), event.Value)
	}
}

func TestExecuteBackupInitSQL(t *testing.T) {
	testCases := []struct {
		name          string
		params        *BackupParams
		setupMysqld   func(*FakeMysqlDaemon)
		wantErr       bool
		wantErrString string
		wantLogMsg    string
		wantNoLogMsg  bool
	}{
		{
			name:         "nil params",
			params:       nil,
			wantErr:      false,
			wantNoLogMsg: true,
		},
		{
			name: "nil InitSQL",
			params: &BackupParams{
				InitSQL: nil,
			},
			wantErr:      false,
			wantNoLogMsg: true,
		},
		{
			name: "no queries",
			params: &BackupParams{
				InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
					Queries: []string{},
				},
			},
			wantErr:      false,
			wantNoLogMsg: true,
		},
		{
			name: "skip queries - tablet type not in allow list",
			params: &BackupParams{
				TabletType: topodatapb.TabletType_PRIMARY,
				InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
					Queries:     []string{"OPTIMIZE TABLE foo"},
					TabletTypes: []topodatapb.TabletType{topodatapb.TabletType_REPLICA, topodatapb.TabletType_RDONLY},
					Timeout:     protoutil.DurationToProto(30 * time.Second),
				},
				Logger: logutil.NewMemoryLogger(),
			},
			wantErr:    false,
			wantLogMsg: "Skipping backup init SQL queries",
		},
		{
			name: "missing timeout",
			params: &BackupParams{
				TabletType: topodatapb.TabletType_REPLICA,
				InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
					Queries:     []string{"OPTIMIZE LOCAL TABLE foo"},
					TabletTypes: []topodatapb.TabletType{topodatapb.TabletType_REPLICA},
					Timeout:     nil,
				},
				Logger: logutil.NewMemoryLogger(),
			},
			wantErr:       true,
			wantErrString: "no timeout provided",
			wantNoLogMsg:  true,
		},
		{
			name: "successful execution with single query",
			params: &BackupParams{
				TabletType: topodatapb.TabletType_PRIMARY,
				InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
					Queries:     []string{"OPTIMIZE LOCAL TABLE foo"},
					TabletTypes: []topodatapb.TabletType{topodatapb.TabletType_PRIMARY, topodatapb.TabletType_REPLICA},
					Timeout:     protoutil.DurationToProto(30 * time.Second),
					FailOnError: true,
				},
				Logger: logutil.NewMemoryLogger(),
			},
			setupMysqld: func(fmd *FakeMysqlDaemon) {
				fmd.ExpectedExecuteSuperQueryList = []string{
					"OPTIMIZE LOCAL TABLE foo",
				}
			},
			wantErr:    false,
			wantLogMsg: "Successfully completed init SQL queries: \"OPTIMIZE LOCAL TABLE foo\"",
		},
		{
			name: "successful execution with multiple queries in order",
			params: &BackupParams{
				TabletType: topodatapb.TabletType_PRIMARY,
				InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
					Queries: []string{
						"SET sql_log_bin=0",
						"SET super_read_only=OFF",
						"SET SESSION innodb_fast_shutdown=0",
						"OPTIMIZE TABLE test.table1",
						"OPTIMIZE TABLE test.table2",
						"SET sql_log_bin=1",
						"SET super_read_only=ON",
					},
					TabletTypes: []topodatapb.TabletType{topodatapb.TabletType_PRIMARY},
					Timeout:     protoutil.DurationToProto(30 * time.Second),
					FailOnError: true,
				},
				Logger: logutil.NewMemoryLogger(),
			},
			setupMysqld: func(fmd *FakeMysqlDaemon) {
				fmd.ExpectedExecuteSuperQueryList = []string{
					"SET sql_log_bin=0",
					"SET super_read_only=OFF",
					"SET SESSION innodb_fast_shutdown=0",
					"OPTIMIZE TABLE test.table1",
					"OPTIMIZE TABLE test.table2",
					"SET sql_log_bin=1",
					"SET super_read_only=ON",
				}
			},
			wantErr:    false,
			wantLogMsg: "Successfully completed init SQL queries",
		},
		{
			name: "no tablet types provided",
			params: &BackupParams{
				TabletType: topodatapb.TabletType_PRIMARY,
				InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
					Queries:     []string{"OPTIMIZE TABLE foo"},
					TabletTypes: []topodatapb.TabletType{},
					Timeout:     protoutil.DurationToProto(30 * time.Second),
				},
				Logger: logutil.NewMemoryLogger(),
			},
			wantErr:       true,
			wantErrString: "backup init SQL queries provided but no tablet types",
			wantNoLogMsg:  true,
		},
		{
			name: "query failure with FailOnError true",
			params: &BackupParams{
				TabletType: topodatapb.TabletType_PRIMARY,
				InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
					Queries:     []string{"INVALID SQL QUERY"},
					TabletTypes: []topodatapb.TabletType{topodatapb.TabletType_PRIMARY},
					Timeout:     protoutil.DurationToProto(30 * time.Second),
					FailOnError: true,
				},
				Logger: logutil.NewMemoryLogger(),
			},
			setupMysqld: func(fmd *FakeMysqlDaemon) {
				fmd.ExpectedExecuteSuperQueryList = []string{"DIFFERENT QUERY"}
			},
			wantErr:       true,
			wantErrString: "failed to execute init SQL queries",
		},
		{
			name: "query failure with FailOnError false",
			params: &BackupParams{
				TabletType: topodatapb.TabletType_PRIMARY,
				InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
					Queries:     []string{"INVALID SQL QUERY"},
					TabletTypes: []topodatapb.TabletType{topodatapb.TabletType_PRIMARY},
					Timeout:     protoutil.DurationToProto(30 * time.Second),
					FailOnError: false,
				},
				Logger: logutil.NewMemoryLogger(),
			},
			setupMysqld: func(fmd *FakeMysqlDaemon) {
				fmd.ExpectedExecuteSuperQueryList = []string{"DIFFERENT QUERY"}
			},
			wantErr:    false,
			wantLogMsg: "Continuing with backup after failed init SQL work",
		},
		{
			name: "zero timeout value",
			params: &BackupParams{
				TabletType: topodatapb.TabletType_PRIMARY,
				InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
					Queries:     []string{"OPTIMIZE TABLE foo"},
					TabletTypes: []topodatapb.TabletType{topodatapb.TabletType_PRIMARY},
					Timeout:     protoutil.DurationToProto(0),
				},
				Logger: logutil.NewMemoryLogger(),
			},
			wantErr:       true,
			wantErrString: "no timeout provided",
			wantNoLogMsg:  true,
		},
	}

	// Test case for context cancellation during query execution.
	t.Run("parent context canceled with query failure", func(t *testing.T) {
		sqldb := fakesqldb.New(t)
		defer sqldb.Close()
		mysqld := NewFakeMysqlDaemon(sqldb)
		defer mysqld.Close()

		// Set up to fail the query.
		mysqld.ExpectedExecuteSuperQueryList = []string{"DIFFERENT QUERY"}
		logger := logutil.NewMemoryLogger()

		params := &BackupParams{
			TabletType: topodatapb.TabletType_PRIMARY,
			InitSQL: &tabletmanagerdatapb.BackupRequest_InitSQL{
				Queries:     []string{"ACTUAL QUERY"},
				TabletTypes: []topodatapb.TabletType{topodatapb.TabletType_PRIMARY},
				Timeout:     protoutil.DurationToProto(30 * time.Second),
				FailOnError: false,
			},
			Mysqld: mysqld,
			Logger: logger,
		}

		ctx, cancel := context.WithCancel(context.Background())
		cancel() // Cancel immediately to simulate parent context cancellation

		err := ExecuteBackupInitSQL(ctx, params)
		require.NoError(t, err) // Should not error because FailOnError is false

		// Check that the appropriate log message was generated.
		found := false
		for _, event := range logger.Events {
			if strings.Contains(event.Value, "Canceling init SQL work due to hitting the configured timeout") {
				found = true
				break
			}
		}
		assert.True(t, found, "Expected log message about canceling init SQL work")
	})

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			var mysqld *FakeMysqlDaemon
			if tc.params != nil && tc.params.InitSQL != nil && len(tc.params.InitSQL.Queries) > 0 && tc.setupMysqld != nil {
				sqldb := fakesqldb.New(t)
				defer sqldb.Close()
				mysqld = NewFakeMysqlDaemon(sqldb)
				defer mysqld.Close()
				tc.setupMysqld(mysqld)
				tc.params.Mysqld = mysqld
			}

			ctx := context.Background()
			err := ExecuteBackupInitSQL(ctx, tc.params)

			if tc.wantErr {
				require.Error(t, err)
				if tc.wantErrString != "" {
					assert.ErrorContains(t, err, tc.wantErrString)
				}
			} else {
				require.NoError(t, err)
			}

			if mysqld != nil && !tc.wantErr {
				assert.NoError(t, mysqld.CheckSuperQueryList())
			}

			// Verify logging too.
			if tc.params != nil && tc.params.Logger != nil {
				memLogger, ok := tc.params.Logger.(*logutil.MemoryLogger)
				require.True(t, ok, "Expected MemoryLogger")

				if tc.wantNoLogMsg {
					assert.Empty(t, memLogger.Events, "Expected no log messages")
				} else if tc.wantLogMsg != "" {
					found := false
					for _, event := range memLogger.Events {
						if strings.Contains(event.Value, tc.wantLogMsg) {
							found = true
							break
						}
					}
					assert.True(t, found, "Expected log message containing %q, got events: %+v", tc.wantLogMsg, memLogger.Events)
				}
			}
		})
	}
}
